package top.hekun.study.agent;

import java.io.File;
import java.lang.instrument.Instrumentation;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.net.URLClassLoader;
import java.security.CodeSource;
import java.security.ProtectionDomain;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * @Author: 贺坤
 * @Date: 2023/4/21 18:11
 */
public class AgentPremain {
    private static final ClassLoader BOOTSTRAP_CLASS_LOADER = null;

    /**
     * 探针入口
     *
     * @param arg
     * @param instrumentation
     */
    public static void premain(String arg, Instrumentation instrumentation) throws Exception {
        System.out.println("探针启动！！！");
        System.out.println("探针传入参数：" + arg);
        // 获取当前jar包
        File archiveFileContains = getArchiveFileContains();
        // 获取路径
        URL urls = archiveFileContains.toURI().toURL();
        // 构造类加载器 与项目进行隔离
        ClassLoader loader = new CompoundableClassLoader(urls);
        // 开始记录
        switchLoggingProperty(loader, () -> {
            // 加载一个类
            loader.loadClass("top.hekun.study.core.agent.AgentPremainDemo")
                    // 找指定方法
                    .getMethod("agentPremainAddTransformer", Instrumentation.class)
                    // 调用方法
                    .invoke(null, instrumentation);
            return null;
        });
    }


    /**
     * 开始记录
     * @param loader 类加载器
     * @param callable 函数
     * @throws Exception
     */
    private static void switchLoggingProperty(ClassLoader loader, Callable<Void> callable)
            throws Exception {
        // 获取当前线程
        final Thread t = Thread.currentThread();
        // 获取线程类加载器
        final ClassLoader ccl = t.getContextClassLoader();
        // 设置线程使用类加载器
        t.setContextClassLoader(loader);
        try {
            // 方法调用
            callable.call();
        } finally {
            // 还原类加载器
            t.setContextClassLoader(ccl);
        }
    }

    /**
     * url 类加加载器
     *
     */
    public static class CompoundableClassLoader extends URLClassLoader {
        private final Set<ClassLoader> externals = new CopyOnWriteArraySet<>();

        CompoundableClassLoader(URL... urls) {
            // super(urls, ClassLoader.getSystemClassLoader());
            super(urls, AgentPremain.BOOTSTRAP_CLASS_LOADER);
        }

        public void add(ClassLoader cl) {
            if (cl != null && !Objects.equals(cl, this)) {
                externals.add(cl);
            }
        }

        @Override
        protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
            try {
                return super.loadClass(name, resolve);
            } catch (ClassNotFoundException e) {
                for (ClassLoader external : externals) {
                    try {
                        final Class<?> aClass = external.loadClass(name);
                        if (resolve) {
                            resolveClass(aClass);
                        }
                        return aClass;
                    } catch (ClassNotFoundException ignore) {
                    }
                }

                throw e;
            }
        }
    }

    /**
     * 获取当前jar包
     * @return
     * @throws URISyntaxException
     */
    private static File getArchiveFileContains() throws URISyntaxException {
        // 获取当前类的域
        final ProtectionDomain protectionDomain = AgentPremain.class.getProtectionDomain();
        // 获取代码源
        final CodeSource codeSource = protectionDomain.getCodeSource();
        // 获取位置URI
        final URI location = (codeSource == null ? null : codeSource.getLocation().toURI());
        // 获取路径
        final String path = (location == null ? null : location.getSchemeSpecificPart());

        if (path == null) {
            throw new IllegalStateException("无法确定代码源存档");
        }
        // 根据路径获取文件
        final File root = new File(path);
        if (!root.exists() || root.isDirectory()) {
            throw new IllegalStateException("无法从中确定代码源存档 " + root);
        }
        return root;
    }

}
